import math
import random

import numpy as np
import os, sys, json, cv2
import torch.utils.data as data
from torch.utils.data import DataLoader

# def get_title(key):
#     return {
#         0: '背景',
#         1: '瓶盖破损',
#         2: '瓶盖变形',
#         3: '瓶盖坏边',
#         4: '瓶盖打旋',
#         5: '瓶盖断点',
#         6: '标贴歪斜',
#         7: '标贴起皱',
#         8: '标贴气泡',
#         9: '喷码正常',
#         10: '喷码异常'
#     }.get(key)
from tqdm import tqdm


def get_title(key):
    return {
        0: '瓶盖破损',
        1: '瓶盖变形',
        2: '瓶盖坏边',
        3: '瓶盖打旋',
        4: '瓶盖断点',
        5: '标贴歪斜',
        6: '标贴起皱',
        7: '标贴气泡',
        8: '喷码正常',
        9: '喷码异常'
    }.get(key)


def loader(path, isResize=True, box=None):
    img = cv2.imread(path)
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

    box[2] = box[0] + box[2]
    box[3] = box[1] + box[3]

    if isResize:
        maxSide = max(img.shape[0], img.shape[1])
        scale = 512.0 / maxSide
        newW = int(img.shape[1] * scale)
        newH = int(img.shape[0] * scale)
        img = cv2.resize(img, (newW, newH))

        box = box * scale

        while img.shape[1] < 512:
            offset = int((512 - img.shape[0]) / 2)
            pad = np.zeros((512, offset, 3), np.uint8)
            img = np.concatenate((pad, img, pad), 1)
            box[0] += offset
            box[2] += offset

        while img.shape[0] < 512:
            offset = int((512 - img.shape[0]) / 2)
            pad = np.zeros((offset, 512, 3), np.uint8)
            img = np.concatenate((pad, img, pad), 0)
            box[1] += offset
            box[3] += offset

    return img, box


def generateItem(all_pairs, root_path, idx):
    labels = np.array(all_pairs[idx][1])
    bbox = np.array(all_pairs[idx][2])
    img, bbox = loader(os.path.join(root_path, all_pairs[idx][0]), isResize=True, box=bbox)

    imgNr = (img - np.mean(img, axis=(0, 1), keepdims=2)).transpose(2, 0, 1)
    # labelIm = cv2.imread("../../labelName.png")
    # cv2.imshow("label", labelIm)
    # img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (255, 0, 255), 2)
    # cv2.imshow(str(all_pairs[idx][1]), img)
    # cv2.waitKey(25)

    return imgNr, bbox, labels, img


class DefectAnno(data.Dataset):
    def __init__(self, file_path,
                 root='/home/leo/Downloads/data_set/industryAss/chongqing1_round1_train1_20191223/images/'):
        self.file_path = file_path
        self.root_path = root
        self.Load()
        random.shuffle(self.all_pairs)
        self.all_pairs = self.all_pairs[0:10]
        pass

    def ReFormatAnno(self):
        self.annotationsReformated = {}
        for anno in self.annotations:
            self.annotationsReformated[anno['image_id']] = anno

    def Load(self):
        dictAnno = json.load(open(self.file_path, "r"))

        self.images = dictAnno['images']  # {file_name:xxx, id: 1, height: xxx, width:xxx}
        self.annotations = dictAnno[
            'annotations']  # {area: 2000, iscrowd: 0 , image_id: 1, bbox:[x,y,w,h], 'category_id':2, 'id': 1 }
        # TODO: annotation exist another id , with should be standard.
        self.categories = dictAnno['categories']  # { supercategory: 'po shun', id: 1 , name: 'p shun'}
        self.ReFormatAnno()

        # " file name , categoryId , bbox
        self.all_pairs = []
        size = len(self.images)
        for i, image in enumerate(self.images):
            if int(image['width']) != 658 or int(image['height']) != 492:
                continue
            if self.annotationsReformated[image['id']]['category_id'] == 0:
                continue
            self.all_pairs.append([image['file_name'],
                                   self.annotationsReformated[image['id']]['category_id'] - 1,
                                   self.annotationsReformated[image['id']]['bbox']]
                                  )

            assert (self.annotationsReformated[image['id']]['category_id'] != None)
            assert (self.annotationsReformated[image['id']]['bbox'] != None)

            print(" appending image {} {} {} {}/{}".format(image['file_name'],
                                                           self.annotationsReformated[image['id']]['category_id'],
                                                           self.annotationsReformated[image['id']]['bbox'],
                                                           i, size))
            assert (self.annotationsReformated[image['id']]['image_id'] == image['id'])

        print(" load complete size {}".format(len(self.all_pairs)))
        return

    def __len__(self):
        return len(self.all_pairs)

    def __getitem__(self, idx):
        return generateItem(self.all_pairs, self.root_path, idx)

    def SplitDataSet(self):
        ratio = 0.9
        annoTrain = DefectAnnoTrain(self, ratio=ratio)
        annoEval = DefectAnnoEval(self, ratio=ratio)
        return annoTrain, annoEval


class DefectAnnoTrain(data.Dataset):
    def __init__(self, DefectAnno, ratio=0.9):
        self.root_path = DefectAnno.root_path
        full_size = len(DefectAnno.all_pairs)
        idx = math.ceil(full_size * ratio)
        self.all_pairs = DefectAnno.all_pairs[:idx]
        print(" Train Set size:{}".format(len(self.all_pairs)))

    def __len__(self):
        return len(self.all_pairs)

    def __getitem__(self, idx):
        return generateItem(self.all_pairs, self.root_path, idx)


class DefectAnnoEval(data.Dataset):
    def __init__(self, DefectAnno, ratio=0.9):
        self.root_path = DefectAnno.root_path
        full_size = len(DefectAnno.all_pairs)
        idx = math.ceil(full_size * ratio)
        self.all_pairs = DefectAnno.all_pairs[idx:]
        print(" Eval Set size:{}".format(len(self.all_pairs)))

    def __len__(self):
        return len(self.all_pairs)

    def __getitem__(self, idx):
        return generateItem(self.all_pairs, self.root_path, idx)


if __name__ == '__main__':
    annosDefectSet = DefectAnno("./annotations.json")

    annoTrain, annoEval = annosDefectSet.SplitDataSet()

    train_dataloader = DataLoader(annoTrain,
                                  batch_size=1,
                                  num_workers=1,
                                  shuffle=True,
                                  pin_memory=False)

    eval_dataloader = DataLoader(annoEval,
                                 batch_size=1,
                                 num_workers=1,
                                 shuffle=True,
                                 pin_memory=False)

    for (images, bbox, label) in tqdm(train_dataloader):
        # print("train literate")
        pass

    for (images, bbox, label) in tqdm(eval_dataloader):
        # print("train literate")
        pass
